{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "829b6f5e-b8d1-4af4-b1b5-2e876c0a95bb",
   "metadata": {},
   "source": [
    "## 特征提取"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d588c49-5939-480d-b1db-8a23d8b26b5d",
   "metadata": {},
   "source": [
    "依赖与环境"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "220577e8-2fed-4b27-817f-f110f528c3d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] DEVICE(5203,e7ffc65c2020,python):2025-08-22-17:00:49.994.794 [mindspore/ccsrc/utils/dlopen_macro.h:165] DlsymAscend] Dynamically load symbol aclprofGetSupportedFeaturesV2 failed, result = /usr/local/Ascend/ascend-toolkit/latest/lib64/libmsprofiler.so: undefined symbol: aclprofGetSupportedFeaturesV2\n",
      "[WARNING] DEVICE(5203,e7ffc65c2020,python):2025-08-22-17:00:49.994.921 [mindspore/ccsrc/utils/dlopen_macro.h:165] DlsymAscend] Dynamically load symbol aclrtEventGetTimestamp failed, result = /usr/local/Ascend/ascend-toolkit/latest/lib64/libascendcl.so: undefined symbol: aclrtEventGetTimestamp\n",
      "/usr/local/miniconda3/lib/python3.9/site-packages/numpy/core/getlimits.py:499: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/usr/local/miniconda3/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n",
      "/usr/local/miniconda3/lib/python3.9/site-packages/numpy/core/getlimits.py:499: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/usr/local/miniconda3/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n",
      "[WARNING] ME(5203:255085730603040,MainProcess):2025-08-22-17:01:02.404.93 [mindspore/context.py:1402] For 'context.set_context', the parameter 'ascend_config' will be deprecated and removed in a future version. Please use the api mindspore.device_context.ascend.op_precision.precision_mode(),\n",
      "                                                       mindspore.device_context.ascend.op_precision.op_precision_mode(),\n",
      "                                                       mindspore.device_context.ascend.op_precision.matmul_allow_hf32(),\n",
      "                                                       mindspore.device_context.ascend.op_precision.conv_allow_hf32(),\n",
      "                                                       mindspore.device_context.ascend.op_tuning.op_compile() instead.\n",
      "Building prefix dict from the default dictionary ...\n",
      "Dumping model to file cache /tmp/jieba.cache\n",
      "Loading model cost 2.132 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mindnlp==? (要求 0.4.1)\n",
      "numpy 已安装\n"
     ]
    }
   ],
   "source": [
    "import sys, platform, os, time, importlib, subprocess\n",
    "import mindspore as ms\n",
    "from mindspore import context\n",
    "\n",
    "def ensure(pkg, ver=None):\n",
    "    try:\n",
    "        m = importlib.import_module(pkg)\n",
    "        if ver:\n",
    "            print(f\"{pkg}=={getattr(m,'__version__','?')} (要求 {ver})\")\n",
    "        else:\n",
    "            print(f\"{pkg} 已安装\")\n",
    "    except Exception as e:\n",
    "        print(f\"缺失 {pkg}，开始安装……\")\n",
    "        cmd = [sys.executable, \"-m\", \"pip\", \"install\", pkg + (f\"=={ver}\" if ver else \"\")]\n",
    "        print(\">>>\", \" \".join(cmd))\n",
    "        subprocess.check_call(cmd)\n",
    "\n",
    "ensure(\"mindnlp\", \"0.4.1\")\n",
    "ensure(\"numpy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5904ef4-0842-49ba-895c-4828e9b30a83",
   "metadata": {},
   "source": [
    "加载 BGE 模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a123d16a-8292-452a-9791-eb5ba2baa16c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading: BAAI/bge-small-zh-v1.5\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BertModel(\n",
       "  (embeddings): BertEmbeddings(\n",
       "    (word_embeddings): Embedding(21128, 512, padding_idx=0)\n",
       "    (position_embeddings): Embedding(512, 512)\n",
       "    (token_type_embeddings): Embedding(2, 512)\n",
       "    (LayerNorm): LayerNorm((512,), eps=1e-12, elementwise_affine=True)\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (encoder): BertEncoder(\n",
       "    (layer): ModuleList(\n",
       "      (0-3): 4 x BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear (512 -> 512)\n",
       "            (key): Linear (512 -> 512)\n",
       "            (value): Linear (512 -> 512)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear (512 -> 512)\n",
       "            (LayerNorm): LayerNorm((512,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear (512 -> 2048)\n",
       "          (intermediate_act_fn): GELU(approximate='none')\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear (2048 -> 512)\n",
       "          (LayerNorm): LayerNorm((512,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (pooler): BertPooler(\n",
       "    (dense): Linear (512 -> 512)\n",
       "    (activation): Tanh()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from mindnlp.transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "MODEL_ID = \"BAAI/bge-small-zh-v1.5\"  \n",
    "print(\"Loading:\", MODEL_ID)\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)   \n",
    "model = AutoModel.from_pretrained(MODEL_ID)\n",
    "model.set_train(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "364d32b0-9e99-4d1c-bd91-8792add0ecd6",
   "metadata": {},
   "source": [
    "文本特征提取的编码函数\n",
    "\n",
    "步骤如下：\n",
    "\n",
    "（1）输入：文本列表 → tokenizer → token 张量\n",
    "\n",
    "（2）模型输出：每个 token 的向量 [B, L, H]（B = 句子数，L = 每句子 token 数，H = 向量维度）\n",
    "\n",
    "（3）平均池化：每句子的向量 = 所有 token 向量的平均值（求和 ÷ token 数 → [B, H]）\n",
    "\n",
    "（4）归一化：把每个向量缩放到长度为1，使得点积=余弦相似度\n",
    "\n",
    "（5）输出：每句子一个固定维度向量 [N, H]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "15892e7f-301d-4ad0-ab76-7d1f5b7f4f44",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import mindspore.ops as ops\n",
    "\n",
    "def encode_texts(texts, max_len=256, batch_size=16):\n",
    "    all_vecs = []\n",
    "    for i in range(0, len(texts), batch_size):\n",
    "        batch = texts[i:i+batch_size]\n",
    "        enc = tokenizer(batch, padding=True, truncation=True, max_length=max_len, return_tensors=\"ms\")\n",
    "        out = model(**enc)\n",
    "        mask = enc[\"attention_mask\"].astype(ms.float32).expand_dims(-1) \n",
    "        summed = ops.sum(out.last_hidden_state * mask, dim=1, keepdim=False)  \n",
    "        counts = ops.clip_by_value(ops.sum(mask, dim=1, keepdim=False), 1.0, 1e9)\n",
    "        mean_pooled = summed / counts  \n",
    "        all_vecs.append(mean_pooled.asnumpy())\n",
    "    X = np.vstack(all_vecs).astype(\"float32\")\n",
    "    X /= (np.linalg.norm(X, axis=1, keepdims=True) + 1e-12)\n",
    "    return X"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32f872d2-75bd-48fa-bac2-047e7ebcbe17",
   "metadata": {},
   "source": [
    "准备语料"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "76b9ad60-dbb6-4ef2-bc8b-508c23042ed9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "语料向量形状: (8, 512)\n"
     ]
    }
   ],
   "source": [
    "corpus = [\n",
    "    \"我喜欢在假期去海边旅行，看看大海、晒晒太阳。\",\n",
    "    \"家里新来了一只可爱的猫咪，特别黏人。\",\n",
    "    \"最近在学深度学习，准备做一个图像检索的小项目。\",\n",
    "    \"今天读到一句话：读书破万卷，下笔如有神。\",\n",
    "    \"科学技术是第一生产力。\",\n",
    "    \"周末打算去爬山，呼吸新鲜空气，顺便拍点照片。\",\n",
    "    \"宠物狗很听话，已经学会了坐下和握手。\",\n",
    "    \"论文要交了，我在写实验部分和结果分析。\",\n",
    "]\n",
    "corpus_embeds = encode_texts(corpus)\n",
    "print(\"语料向量形状:\", corpus_embeds.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ca03df1-5af2-43f4-87c8-76da8aafc4ea",
   "metadata": {},
   "source": [
    "语义检索：给定查询，找到最相似的句子 Top-K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "46a6d31d-cbd9-4219-8191-1b5f574a2708",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_topk(query_texts, index_texts, index_embeds, topk=3):\n",
    "    q_emb = encode_texts(query_texts)\n",
    "    sims = q_emb @ index_embeds.T    \n",
    "    results = []\n",
    "    for qi, q in enumerate(query_texts):\n",
    "        k = min(topk, sims.shape[1])\n",
    "        cand = np.argpartition(-sims[qi], range(k))[:k]\n",
    "        cand = cand[np.argsort(-sims[qi, cand])]  \n",
    "        items = [(int(j), float(sims[qi, j]), index_texts[j]) for j in cand]\n",
    "        results.append({\"query\": q, \"topk\": items})\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2bc45bc8-452f-4080-9a6b-50137dc12931",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== 语义检索结果（Top-3） ===\n",
      "\n",
      "Query: 找关于旅游出行的句子\n",
      "  1. score=0.4897 | #0 | 我喜欢在假期去海边旅行，看看大海、晒晒太阳。\n",
      "  2. score=0.3406 | #5 | 周末打算去爬山，呼吸新鲜空气，顺便拍点照片。\n",
      "  3. score=0.2899 | #2 | 最近在学深度学习，准备做一个图像检索的小项目。\n",
      "\n",
      "Query: 和科研学习相关的内容\n",
      "  1. score=0.5057 | #2 | 最近在学深度学习，准备做一个图像检索的小项目。\n",
      "  2. score=0.4604 | #4 | 科学技术是第一生产力。\n",
      "  3. score=0.4596 | #7 | 论文要交了，我在写实验部分和结果分析。\n",
      "\n",
      "Query: 与宠物有关\n",
      "  1. score=0.4877 | #6 | 宠物狗很听话，已经学会了坐下和握手。\n",
      "  2. score=0.4686 | #1 | 家里新来了一只可爱的猫咪，特别黏人。\n",
      "  3. score=0.4074 | #2 | 最近在学深度学习，准备做一个图像检索的小项目。\n"
     ]
    }
   ],
   "source": [
    "queries = [\n",
    "    \"找关于旅游出行的句子\",\n",
    "    \"和科研学习相关的内容\",\n",
    "    \"与宠物有关\"\n",
    "]\n",
    "\n",
    "res = search_topk(queries, corpus, corpus_embeds, topk=3)\n",
    "\n",
    "print(\"\\n=== 语义检索结果（Top-3） ===\")\n",
    "for r in res:\n",
    "    print(f\"\\nQuery: {r['query']}\")\n",
    "    for rank, (idx, s, txt) in enumerate(r[\"topk\"], 1):\n",
    "        print(f\"  {rank}. score={s:.4f} | #{idx} | {txt}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
